#!/usr/bin/python
#coding:utf-8

# Copyright 2011 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under theglim Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
## Open a (pinhole camera model) picture, find its orientation and
## extract the rectified.
##
## Changes that must be performed as soon as possible: read the
## intrinsic parameters from somewhere, and make it easy to switch to
## the equirectangular model.

import sys

import time

import matplotlib

if __name__ == '__main__':
    if sys.argv[0][-7:] == '-nox.py':
        matplotlib.use('Agg') 

import plot_aux

from pylab import *

import scipy.io

import scipy.signal

from numpy import dot

# from camori import PicturePinhole, PictureEqrec, quaternion_to_matrix, measure_error, quaternion_product, dir_colors
from camori_harris import PictureHarris
from camori_aux import *
from quaternion import Quat, random_quaternion

import simplejson

import scipy.optimize
#from fmin_sa import Sampling
from simplex_sa import SimplexSO3

import Image

from plot_aux import *

set_printoptions(precision=7)



import filterSQP

## Conic constraint
def val_c(x):
    return np.linalg.norm(x)
def grad_c(x):
    return x/np.linalg.norm(x)
def hess_c(x):
    nx = np.linalg.norm(x)
    return (nx**2 * np.identity(x.shape[0]) - np.outer(x,x)) / nx**3

## Target function
def val_f(x,*fargs):
    return harris_angle_error(fargs[0],fargs[1],fargs[2],fargs[3],
                              x, fargs[4],fargs[5],fargs[6],fargs[7])
def grad_f(x,*fargs):
    return harris_angle_error_gradient(fargs[0],fargs[1],fargs[2],fargs[3],
                              x, fargs[4],fargs[5],fargs[6],fargs[7])
def hess_f(x,*fargs):
    return harris_angle_error_hessian(fargs[0],fargs[1],fargs[2],fargs[3],
                              x, fargs[4],fargs[5],fargs[6],fargs[7])

def harris_map(ox, oy, kappa_in, p):
    # cdef double fn = 1.0 / sqrt(fabs(1 + kappa * (px[0] * px[0] + py[0] * py[0])))
    p = p - array([ox,oy])
    kappa = 1e-9 * kappa_in
    fn = np.abs(1 + kappa * (p[0]*p[0] + p[1]*p[1]))**-0.5
    return p*fn
    
def aligned_quaternion(v):
    ## The largest component
    ll = argmax(np.abs(v))
    ee = zeros(3)
    ee[ll] = sign(v[ll])
    q=Quat(cross(v, ee))
    R = q.sqrt()

    if sign(v[ll]) > 0:
        if ll == 1:
            R = R*Quat(1,-1,-1,-1).normalize()
        if ll == 2:
            R = R*Quat(1,1,1,1).normalize()

    if sign(v[ll]) < 0:
        if ll == 0:
            R = R*Quat(0,0,1,0).normalize()
        if ll == 1:
            R = R*Quat(1,1,-1,1).normalize()
        if ll == 2:
            R = R*Quat(1,-1,-1,1).normalize()
    return R.inverse()

if __name__ == '__main__':

    rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    PlotStuff=True
    if sys.argv[0][-7:] == '-nop.py':
        PlotStuff=False

    if PlotStuff:
        ## Plot stuff immediately
        ion()
        set_printoptions(precision=3)

    #################################################################
    ## Load image and initialize pic object

    ## Sets filename from input argument
    if len(sys.argv) < 3:
        print sys.argv[0], '<job_file.json> <frame_number>'
        raise Exception('Insufficient number of parameters')

    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    fileroot = job_params['root_directory']

    framenum = int(sys.argv[2])
    filename = fileroot+'/frames/'+job_params['filename_format']%framenum

    im = Image.open(filename)
    frame = array(im.convert('RGB'), dtype=float)
    imr = array(im.convert('RGB'), dtype=float)
    imr = imr[:,:,:3] #remove alpha channel

    # Smooth out
    if ("gaussian_smoothing_factor" in job_params.keys() and 
        job_params["gaussian_smoothing_factor"] > 0):
        for c in range(3):
            imr[:,:,c] = scipy.ndimage.gaussian_filter( imr[:,:,c], \
                          double(job_params["gaussian_smoothing_factor"]))

    # scharr_filter=scipy.array([
    #     [ -3, 0, 3],
    #     [-10, 0,10],
    #     [ -3, 0, 3] ])
    shigeru_filter=-scipy.array([
            [ -0.003776, -0.010199, 0., 0.010199, 0.003776 ],
            [ -0.026786, -0.070844, 0., 0.070844, 0.026786 ],
            [ -0.046548, -0.122572, 0., 0.122572, 0.046548 ],
            [ -0.026786, -0.070844, 0., 0.070844, 0.026786 ],
            [ -0.003776, -0.010199, 0., 0.010199, 0.003776 ]
            ])

    
    gradx = scipy.zeros(imr.shape, dtype=np.float64)
    grady = scipy.zeros(imr.shape, dtype=np.float64)
    gradxx = scipy.zeros(imr.shape, dtype=np.float64)
    gradxy = scipy.zeros(imr.shape, dtype=np.float64)
    gradyy = scipy.zeros(imr.shape, dtype=np.float64)

    ttini = time.time()
    the_filter = shigeru_filter
    for c in range(3):
        scipy.ndimage.convolve(imr[:,:,c], the_filter, gradx[:,:,c])
        scipy.ndimage.convolve(imr[:,:,c], the_filter.T, grady[:,:,c])
    for c in range(3):
        scipy.ndimage.convolve(gradx[:,:,c], the_filter, gradxx[:,:,c])
        scipy.ndimage.convolve(gradx[:,:,c], the_filter.T, gradxy[:,:,c])
        scipy.ndimage.convolve(grady[:,:,c], the_filter.T, gradyy[:,:,c])
    print 'convs', time.time() - ttini


    gstep = job_params['grid_spacing']
    glim = job_params['gradient_threshold']

    ttini = time.time()
    edgels = cython_edgel_extractor(gstep, glim**2, gradx, grady, gradxx, gradxy, gradyy)
    print 'new edgel extractor', time.time() - ttini



    ion()

    # figure(2)
    # imshow(lam_p, vmin=-1e4, vmax=1e4)

    figure(6)
    imshow(imr/256.)
    vv = 2*array([-1.0,1.0])
    for px,py,ux,uy in edgels:
        plot(px+uy*vv, py-ux*vv, 'r-')

